#!/usr/bin/env python  
import rospy  
from sensor_msgs.msg import Imu  
from nav_msgs.msg import Odometry  
from geometry_msgs.msg import Quaternion  
import tf  
import numpy as np  
  
class ImuToOdom:  
    def __init__(self):  
        # 初始化节点  
        rospy.init_node('imu_to_odom', anonymous=True)  
  
        # 订阅IMU数据  
        self.imu_sub = rospy.Subscriber('/camera/imu', Imu, self.imu_callback, queue_size=10)  
  
        # 发布Odometry数据  
        self.odom_pub = rospy.Publisher('/odom', Odometry, queue_size=10)  
  
        # 初始化Odometry消息  
        self.odom = Odometry()  
        self.odom.header.frame_id = "odom"  
        self.odom.child_frame_id = "base_link"
  
        # 初始化位置和姿态  
        self.position = np.array([0.0, 0.0, 0.0])  
        self.orientation_euler = np.array([0.0, 0.0, 0.0])  # 初始欧拉角  
        self.orientation_q = Quaternion(*tf.transformations.quaternion_from_euler(0, 0, 0))  
  
        # 积分参数  
        self.last_time = rospy.Time.now()  
        self.last_orientation_euler = self.orientation_euler.copy()  
  
        # 设置ROS循环率  
        self.rate = rospy.Rate(50)  # 50Hz  
  
    def imu_callback(self, imu_msg):  
        # 计算时间差  
        dt = (imu_msg.header.stamp - self.last_time).to_sec()  
        self.last_time = imu_msg.header.stamp  
  
        # 欧拉积分来更新方向（这里仅使用Z轴角速度）  
        # 注意：这种方法在长时间后会产生很大的误差  
        angular_velocity_z = imu_msg.angular_velocity.z  
        delta_theta = angular_velocity_z * dt  
  
        # 更新欧拉角  
        self.orientation_euler[2] += delta_theta  
  
        # 将欧拉角转换为四元数  
        self.orientation_q = Quaternion(*tf.transformations.quaternion_from_euler(  
            self.orientation_euler[0], self.orientation_euler[1], self.orientation_euler[2]))  
  
        # 设置Odometry消息  
        self.odom.header.stamp = imu_msg.header.stamp  
        self.odom.header.frame_id = "odom"  
        self.odom.pose.pose.position.x = self.position[0]  
        self.odom.pose.pose.position.y = self.position[1]  
        self.odom.pose.pose.position.z = self.position[2]  
        self.odom.pose.pose.orientation = self.orientation_q  
  
        # 假设没有线速度  
        self.odom.twist.twist.linear.x = 0.0  
        self.odom.twist.twist.linear.y = 0.0  
        self.odom.twist.twist.linear.z = 0.0  
  
        # 角速度可以直接从IMU消息中获取，但这里我们只更新了方向  
  
        # 发布Odometry数据  
        self.odom_pub.publish(self.odom)  
  
        # 循环等待  
        self.rate.sleep()  
  
if __name__ == '__main__':  
    try:  
        imu_to_odom = ImuToOdom()  
        rospy.spin()  
    except rospy.ROSInterruptException:  
        pass